import networkx as nx
import numpy as np

# Find line structures (paths of specific lengths)
def find_line_structures(graph, length=3):
    """
    Find all line structures (paths) of specified length in the graph.
    
    Args:
        graph: NetworkX graph representing the coupling map
        length: Length of the path to find
        
    Returns:
        List of paths (each path is a list of node indices following coupling order)
    """
    line_structures = []
    
    # Helper function to perform DFS to find paths of specific length
    def dfs_paths(current_node, current_path, path_length):
        # If we've reached the desired path length, add the path to our list
        if len(current_path) == path_length:
            line_structures.append(current_path.copy())
            return
        
        # Explore neighbors
        for neighbor in graph.neighbors(current_node):
            # Skip nodes we've already visited in this path
            if neighbor not in current_path:
                current_path.append(neighbor)
                dfs_paths(neighbor, current_path, path_length)
                current_path.pop()  # Backtrack
    
    # Start DFS from each node in the graph
    for node in graph.nodes():
        dfs_paths(node, [node], length)
    
    # Remove duplicates (paths that are the same but in reverse order)
    unique_lines = []
    for line in line_structures:
        # Convert to tuple for easier comparison
        line_tuple = tuple(line)
        line_reversed = tuple(reversed(line))
        
        # Check if this line or its reverse is already in unique_lines
        if not any(line_tuple == tuple(existing) or line_reversed == tuple(existing) 
                  for existing in unique_lines):
            unique_lines.append(line)
    
    return unique_lines

# Find grid structures (2x2 grid or larger)
def find_grid_structures(graph, min_size=2):
    """
    Find grid-like structures in the graph following coupling order.
    
    Args:
        graph: NetworkX graph representing the coupling map
        min_size: Minimum size of the grid (e.g., 2 for a 2x2 grid)
        
    Returns:
        List of grid structures (each grid is a list of node indices in row-major order)
    """
    grid_structures = []
    
    # Find 2x2 grid structures first (basic building blocks)
    basic_grids = []
    for node in graph.nodes():
        neighbors = list(graph.neighbors(node))
        for i, n1 in enumerate(neighbors):
            for n2 in neighbors[i+1:]:
                # Check if n1 and n2 are connected to a common node (other than the original node)
                n1_neighbors = set(graph.neighbors(n1))
                n2_neighbors = set(graph.neighbors(n2))
                common_neighbors = n1_neighbors.intersection(n2_neighbors) - {node}
                
                for common in common_neighbors:
                    # Arrange in row-major order based on node indices
                    grid_nodes = [node, n1, n2, common]
                    # Sort to ensure consistent ordering
                    grid = sorted(grid_nodes)
                    if grid not in basic_grids:
                        basic_grids.append(grid)
    
    # Add all 2x2 grids to the results
    grid_structures.extend(basic_grids)
    
    # For larger grids (min_size > 2)
    if min_size > 2:
        # Try to find rectangular grids of size min_size x min_size or larger
        for size in range(min_size, min_size + 1):
            # For each potential grid of size x size
            for node_set in nx.connected_components(graph):
                if len(node_set) >= size * size:
                    subgraph = graph.subgraph(node_set)
                    # Check if this subgraph could form a grid
                    # This is a simplified approach - in a real implementation,
                    # you would verify the grid connectivity pattern more rigorously
                    if all(2 <= subgraph.degree(n) <= 4 for n in subgraph.nodes()):
                        # Try to arrange nodes in a grid pattern
                        # For simplicity, we're just taking the first size*size nodes
                        potential_grid = sorted(list(subgraph.nodes())[:size*size])
                        if len(potential_grid) == size * size:
                            grid_structures.append(potential_grid)
    
    # Remove duplicates while preserving order
    unique_grids = []
    for grid in grid_structures:
        grid_tuple = tuple(grid)
        if grid_tuple not in [tuple(g) for g in unique_grids]:
            unique_grids.append(grid)
    
    return unique_grids

def find_knn_qubits(graph, k=3, source_qubit=None, exact_distance=False):
    """
    Find qubits within a Manhattan distance of k from each qubit in the coupling graph,
    or from a specific source qubit if provided.
    
    Args:
        graph (networkx.Graph): The graph representation of the coupling map
        k (int): The Manhattan distance to consider (path length in the graph)
        source_qubit (int, optional): Specific qubit to find neighbors for within distance k.
                                     If None, finds neighbors for all qubits.
        exact_distance (bool): If True, only return qubits at exactly distance k.
                              If False, return qubits at distance <= k.
    
    Returns:
        dict: A dictionary where keys are qubit indices and values are lists of 
              qubits that are at/within Manhattan distance k (sorted by distance and then by qubit index).
              If source_qubit is specified, returns a dictionary with only one entry.
    """
    knn_dict = {}
    
    # If source_qubit is specified, only process that qubit
    nodes_to_process = [source_qubit] if source_qubit is not None else graph.nodes()
    
    for node in nodes_to_process:
        # Calculate shortest path lengths from this node to all other nodes
        path_lengths = nx.single_source_shortest_path_length(graph, node)
        
        # Filter nodes based on Manhattan distance
        if exact_distance:
            # Only include nodes at exactly distance k
            neighbors = [
                target for target, length in path_lengths.items() 
                if target != node and length == k
            ]
        else:
            # Include nodes within distance k (less than or equal to k)
            neighbors = [
                target for target, length in path_lengths.items() 
                if target != node and length <= k
            ]
        
        # Sort by distance first, then by node index for consistent ordering
        sorted_neighbors = sorted(
            [(target, path_lengths[target]) for target in neighbors],
            key=lambda x: (x[1], x[0])
        )
        
        # Extract just the node indices
        knn_dict[node] = [neighbor for neighbor, _ in sorted_neighbors]
    
    return knn_dict


if __name__ == "__main__":
    from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
    from qiskit_ibm_runtime.fake_provider import FakeGuadalupeV2

    backend_name = "fake_guadalupe"
    backend = FakeGuadalupeV2()
    
    # Analyze the coupling map to identify specific qubit structures (line, grid, etc.)
    coupling_map = backend.configuration().coupling_map
    # print(f"Coupling map: {coupling_map}")

    # Create a graph representation of the coupling map
    import networkx as nx
    import matplotlib.pyplot as plt

    # Create a graph from the coupling map
    coupling_graph = nx.Graph()
    for i, j in coupling_map:
        coupling_graph.add_edge(i, j)

    # Visualize the coupling graph
    # plt.figure(figsize=(10, 8))
    # pos = nx.spring_layout(coupling_graph, seed=42)  # For reproducible layout
    # nx.draw(coupling_graph, pos, with_labels=True, node_color='lightblue', 
    #         node_size=500, font_weight='bold', font_size=12, edge_color='gray')
    # plt.title(f"Coupling Map for {backend_name}")
    # plt.savefig(f"./data/coupling_map_{backend_name}.png", dpi=300, bbox_inches='tight')
    # plt.show()

    # Find line structures of length 3 and 4
    line_structures_3 = find_line_structures(coupling_graph, length=3)
    line_structures_4 = find_line_structures(coupling_graph, length=4)

    # Find grid structures
    grid_structures_2 = find_grid_structures(coupling_graph, min_size=2)
    grid_structures_3 = find_grid_structures(coupling_graph, min_size=3)

    # Find k-nearest neighbor (kNN) qubit indices for each qubit in the coupling graph
    knn_dict_3 = find_knn_qubits(coupling_graph, k=3)
    knn_dict_4 = find_knn_qubits(coupling_graph, k=4)
    knn_dict_4_source_8 = find_knn_qubits(coupling_graph, k=4, source_qubit=8, exact_distance=True)

    print(f"Found {len(line_structures_3)} line structures of length 3")
    print(f"Found {len(line_structures_4)} line structures of length 4")
    print(f"Found {len(grid_structures_2)} grid structures (2x2)")
    print(f"Found {len(grid_structures_3)} grid structures (3x3)")
    print(f"Found {len(knn_dict_3)} kNN qubit indices for each qubit (k=3)")
    print(f"Found {len(knn_dict_4)} kNN qubit indices for each qubit (k=4)")
    print(f"Found {len(knn_dict_4_source_8)} kNN qubit indices for qubit 8 (k=4)")

    structures = {
        'line_3': line_structures_3,
        'line_4': line_structures_4,
        'grid_2x2': grid_structures_2,
        'grid_3x3': grid_structures_3,
        'knn_3': knn_dict_3,
        'knn_4': knn_dict_4,
        'knn_4_source_8': knn_dict_4_source_8
    }

    print(structures)